#include <thread>
#include <chrono>
#include <algorithm>
#include <iostream>

#include "SocketIOManager.hpp"

namespace cybertron {
	SocketIOManager::SocketIOManager()
	{
#ifdef IO_MULTIPLEXING_MODE_EPOLL
		std::cout << "epoll mode enabled\n";
		_epoll_fd = epoll_create1(0);
		if (_epoll_fd < 0) {
			std::cout << "epoll_create error\n";
			exit(-1);
		}
		_epoll_events = (epoll_event *)calloc(MAX_EPOLL_FDS, sizeof(struct epoll_event));
		if (_epoll_events == NULL) {
			std::cout << ("Unable to allocate memory for epoll_events\n");
			exit(-1);
		}
#else
		FD_ZERO(&_rfds);
		FD_ZERO(&_wfds);
		_maxfd = 0;
#endif
		_io_set_size = 0;
		_last_check_timeout = ::time(NULL);
	}

	SocketIOManager::~SocketIOManager()
	{
#ifdef IO_MULTIPLEXING_MODE_EPOLL
		close(_epoll_fd);
		free(_epoll_events);
#endif
		for (auto io : _io_set) {
			delete io;
		}
	}

	void SocketIOManager::add_io(SocketIO * s)
	{
		_io_set.insert(s);
		_io_set_size = _io_set.size();
#ifdef IO_MULTIPLEXING_MODE_EPOLL
		struct epoll_event event = { 0 };
		event.data.ptr = s;
		if (s->has_read_flag()) {
			event.events |= EPOLLIN;
		}
		if (s->has_write_flag()) {
			event.events |= EPOLLOUT;
		}

		if (epoll_ctl(_epoll_fd, EPOLL_CTL_ADD, s->get_socket(), &event) < 0) {
			std::cout << "epoll add fail\n";
		}
#else
		FD_ZERO(&_rfds);
		FD_ZERO(&_wfds);
		_maxfd = 0;

		for (auto fd : _io_set) {
			if (fd->has_read_flag()) {
				FD_SET(fd->get_socket(), &_rfds);
			}
			if (fd->has_write_flag()) {
				FD_SET(fd->get_socket(), &_wfds);
			}
			_maxfd = (std::max)(_maxfd, fd->get_socket());
		}
#endif
	}

	void SocketIOManager::del_io(SocketIO * s)
	{
		_io_set.erase(s);
		_io_set_size = _io_set.size();
#ifdef IO_MULTIPLEXING_MODE_EPOLL
		if (epoll_ctl(_epoll_fd, EPOLL_CTL_DEL, s->get_socket(), NULL) < 0) {
			std::cout << "epoll del fail\n";
		}
#else
		FD_ZERO(&_rfds);
		FD_ZERO(&_wfds);
		_maxfd = 0;

		for (auto fd : _io_set) {
			if (fd->has_read_flag()) {
				FD_SET(fd->get_socket(), &_rfds);
			}
			if (fd->has_write_flag()) {
				FD_SET(fd->get_socket(), &_wfds);
			}
			_maxfd = (std::max)(_maxfd, fd->get_socket());
		}
#endif
		delete s;
	}

	int SocketIOManager::perform_io(std::unordered_set<AcceptSocket *>* accepts, std::unordered_set<SocketIO*>* errors, int timeoutms)
	{
#ifdef IO_MULTIPLEXING_MODE_EPOLL
		int nready = epoll_wait(_epoll_fd, _epoll_events, MAX_EPOLL_FDS, timeoutms);
		for (int i = 0; i < nready; i++) {
			SocketIO *s = (SocketIO *)(_epoll_events[i].data.ptr);
			if (_epoll_events[i].events & EPOLLERR || _epoll_events[i].events & EPOLLHUP) {
				if (errors) {
					errors->insert(s);
				}
			}
			if (_epoll_events[i].events & EPOLLIN) {
				if (s->has_accept_flag()) { //server socket accept a new connection
					AcceptSocket *p = s->on_accept();
					if (!p) {
						continue; //accept return error
					}
					else if (accepts) {
						accepts->insert(p);
					}
					else {  //no interesting in accepted new connections
						SocketHelper::close(p->_socket);
						delete p;
					}
				}
				else {
					int ret = s->on_read();
					if (ret != 0 && errors) {
						errors->insert(s);
					}
				}
			}
			if (_epoll_events[i].events & EPOLLOUT) {
				int ret = s->on_read();
				if (ret != 0 && errors) {
					errors->insert(s);
				}
			}
		}
#else

#ifdef _WIN32
		if (_io_set.empty()) {
			std::this_thread::sleep_for(std::chrono::milliseconds(20));
			return 0;
		}
#endif

		struct timeval timeout;
		timeout.tv_sec = 0;
		timeout.tv_usec = timeoutms * 1000;

		fd_set rfds = _rfds;
		fd_set wfds = _wfds;

		int ret = ::select(static_cast<int>(_maxfd + 1), &rfds, &wfds, NULL, &timeout);

		if (ret <= 0) {
			return 0;
		}

		for (auto s : _io_set) {
			if (FD_ISSET(s->get_socket(), &rfds)) {
				if (s->has_accept_flag()) { //server socket accept a new connection
					AcceptSocket *p = s->on_accept();
					if (!p) {
						continue; //accept return error
					}
					else if (accepts) {
						accepts->insert(p);
					}
					else {  //no interesting in accepted new connections
						SocketHelper::close(p->_socket);
						delete p;
					}
				}
				else {
					int ret = s->on_read();
					if (ret != 0 && errors) {
						errors->insert(s);
					}
				}
			}

			if (FD_ISSET(s->get_socket(), &wfds)) {
				if (s->on_write() != 0 && errors) {
					errors->insert(s);
				}
			}
		}
#endif
		return 0;
	}

	void SocketIOManager::check_timeout()
	{
		auto now = ::time(NULL);
		if (now - _last_check_timeout > 10) {
			for (auto it = _io_set.begin(); it != _io_set.end();) {
				if ((*it)->is_timeout(now, 120)) {
					std::cout << "timeout disconnecting\n";
					//delete (*it);
					//_io_set.erase(it++);
					del_io(*it);
					return; // del_io will modify _io_set, so the loop can't continue
				}
				else {
					it++;
				}
			}
			_last_check_timeout = now;
		}
	}

}